package com.zfei.serviceA.ribbon;

import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.client.loadbalancer.LoadBalancerRequest;
import org.springframework.cloud.client.loadbalancer.ServiceRequestWrapper;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.AsyncClientHttpRequestExecution;
import org.springframework.http.client.AsyncClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.util.concurrent.ListenableFuture;

import java.io.IOException;
import java.net.URI;


public class SelfAsyncLoadBalancerInterceptor implements AsyncClientHttpRequestInterceptor {

    private LoadBalancerClient loadBalancer;

    public SelfAsyncLoadBalancerInterceptor(LoadBalancerClient loadBalancer) {
        this.loadBalancer = loadBalancer;
    }

    @Override
    public ListenableFuture<ClientHttpResponse> intercept(final HttpRequest request, final byte[] body,
                                                          final AsyncClientHttpRequestExecution execution) throws IOException {
        final URI originalUri = request.getURI();
        String serviceName = originalUri.getHost();

        return this.loadBalancer.execute(serviceName,
                new LoadBalancerRequest<ListenableFuture<ClientHttpResponse>>() {
                    @Override
                    public ListenableFuture<ClientHttpResponse> apply(final ServiceInstance instance)
                            throws Exception {
                        HttpRequest serviceRequest = new ServiceRequestWrapper(request,
                                instance, loadBalancer);
                        return execution.executeAsync(serviceRequest, body);
                    }

                });
    }
}